/*
 * Copyright (c) 2010 Toni Spets <toni.spets@iki.fi>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include "common.h"
#include "launcher.h"

struct hook games[][MAX_HOOKS] =
{
    /* Command & Conquer 1.06 */
    {
        { 0x0045BFB2, hook_IPX_Initialise },
        { 0x0045C59E, hook_IPX_Start_Listening95 },
        { 0x0045C5E4, hook_IPX_Shut_Down95 },
        { 0x0045C79F, hook_IPX_Send_Packet95 },
        { 0x0045C7D0, hook_IPX_Broadcast_Packet95 },
        { 0x0045D45D, hook_IPX_Get_Outstanding_Buffer95 },
        { 0x0045DA08, hook_IPX_Open_Socket95 },
        { 0x0045DA18, hook_IPX_Close_Socket95 },
        { 0x0045DA26, hook_IPX_Get_Connection_Number95 },
        { 0x0045DA33, hook_IPX_Broadcast_Packet95 },
        { 0x0045DA45, hook_IPX_Get_Local_Target95 },
        { 0, NULL }
    },
    /* Command & Conquer: Red Alert 3.03 */
    {
        // 2
        { 0x005A8C8C, hook_bind },
        { 0x005A8E45, hook_bind },
        // 3
        { 0x005A8095, hook_closesocket },
        { 0x005A8D51, hook_closesocket },
        // 7
        { 0x005A8CD6, hook_getsockopt },
        // 17
        { 0x005A8FF8, hook_recvfrom },
        // 20
        { 0x005A9207, hook_sendto },
        // 21
        { 0x005A847B, hook_setsockopt },
        { 0x005A84C4, hook_setsockopt },
        { 0x005A8EAD, hook_setsockopt },
        { 0x005A8EF5, hook_setsockopt },
        { 0x005A8F3D, hook_setsockopt },
        // 23
        { 0x005A8C56, hook_socket },
        { 0x005A8DD0, hook_socket },
        // 101
        { 0x005A80CC, hook_WSAAsyncSelect },
        { 0, NULL }
    },
    /* Command & Conquer: Tiberian Sun 1.03 */
    {
        // 2
        { 0x006A037A, hook_bind },
        { 0x006A0749, hook_bind },
        // 3
        { 0x006A0485, hook_closesocket },
        { 0x006A101F, hook_closesocket },
        // 7
        { 0x006A03DA, hook_getsockopt },
        { 0x006A0445, hook_getsockopt },
        // 17
        { 0x006A0A66, hook_recvfrom },
        // 20
        { 0x006A09A6, hook_sendto },
        // 21
        { 0x006A07B7, hook_setsockopt },
        { 0x006A0808, hook_setsockopt },
        { 0x006A0859, hook_setsockopt },
        { 0x006A175B, hook_setsockopt },
        { 0x006A179F, hook_setsockopt },
        // 23
        { 0x006A0341, hook_socket },
        { 0x006A0625, hook_socket },
        // 101
        { 0x006A1049, hook_WSAAsyncSelect },
        { 0, NULL }
    },
    /* Command & Conquer: Red Alert 2 1.006 */
    {
        // 2
        { 0x0076A0FA, hook_bind },
        { 0x0076A54B, hook_bind },
        // 3
        { 0x0076A219, hook_closesocket },
        { 0x0076AE4F, hook_closesocket },
        // 7
        { 0x0076A15A, hook_getsockopt },
        { 0x0076A1C5, hook_getsockopt },
        // 17
        { 0x0076A866, hook_recvfrom },
        // 20
        { 0x0076A7A6, hook_sendto },
        // 21
        { 0x0076A5B9, hook_setsockopt },
        { 0x0076A60A, hook_setsockopt },
        { 0x0076A65B, hook_setsockopt },
        { 0x0076B912, hook_setsockopt },
        { 0x0076B950, hook_setsockopt },
        // 23
        { 0x0076A0C1, hook_socket },
        { 0x0076A427, hook_socket },
        // 101
        { 0x0076AE9B, hook_WSAAsyncSelect },
        { 0x0076AF2E, hook_WSAAsyncSelect },
        { 0, NULL }
    },
    /* Command & Conquer: Red Alert 2 Yuri's Revenge 1.001 */
    {
        // 2
        { 0x007B0E1A, hook_bind },
        { 0x007B126B, hook_bind },
        // 3
        { 0x007B0F39, hook_closesocket },
        { 0x007B1B8F, hook_closesocket },
        // 7
        { 0x007B0E7A, hook_getsockopt },
        { 0x007B0EE5, hook_getsockopt },
        // 17
        { 0x007B1586, hook_recvfrom },
        // 20
        { 0x007B14C6, hook_sendto },
        // 21
        { 0x007B12D9, hook_setsockopt },
        { 0x007B132A, hook_setsockopt },
        { 0x007B137B, hook_setsockopt },
        { 0x007B2652, hook_setsockopt },
        { 0x007B2690, hook_setsockopt },
        // 23
        { 0x007B0DE1, hook_socket },
        { 0x007B1147, hook_socket },
        // 101
        { 0x007B1BDB, hook_WSAAsyncSelect },
        { 0x007B1C6E, hook_WSAAsyncSelect },
        { 0, NULL }
    },
    {
        { 0, NULL }
    }
};

BOOL FileExists(const char *path)
{
    FILE* file;
    if( (file = fopen(path, "r")) )
    {
        fclose(file);
        return TRUE;
    }
    return FALSE;
}

HANDLE hProcess = NULL;
HANDLE hThread = NULL;

// replacement for dirname() POSIX function (also keeps internal copy of the path)
char *GetDirectory(const char *path)
{
    static char buf[MAX_PATH];
    char *ptr;
    strncpy(buf, path, MAX_PATH);
    ptr = strrchr(buf, '\\');
    if(ptr)
    {
        *(ptr+1) = 0;
        return buf;
    }

    return NULL;
}

char *GetFile(const char *path)
{
    static char buf[MAX_PATH];
    char *ptr;
    strncpy(buf, path, MAX_PATH);
    ptr = strrchr(buf, '\\');
    if(ptr)
    {
        return (ptr+1);
    }

    return buf;
}

int cncnet_launch(const char *gameExe, const char *gameParams)
{
    PROCESS_INFORMATION pInfo;
    STARTUPINFOA sInfo;

    struct hook (*game)[MAX_HOOKS] = NULL;

    // Necessary for a second launch
    game = games;
    hProcess = NULL;
    hThread = NULL;

    char *gamePath = NULL;
    char gameParamsFull[MAX_PATH];

    DEBUG_EVENT ev;
    CONTEXT ctx;
    char bp[5] = { 0xCC, 0x90, 0x90, 0x90, 0x90 };
    struct hook *hk;
    DWORD dwWritten;
    unsigned char byte;
    BOOL match = FALSE;
    DWORD dwProcessId = 0;
    #if _DEBUG
    char *ptr;
    #endif
    int processes = 0;

    if(gameExe == NULL)
    {
        return CNCNET_ERROR_FILE_NOT_FOUND;
    }

    if(!FileExists(gameExe))
    {
        return CNCNET_ERROR_FILE_NOT_FOUND;
    }

    ZeroMemory(&sInfo, sizeof(STARTUPINFO));
    sInfo.cb = sizeof(sInfo);
    ZeroMemory(&pInfo, sizeof(PROCESS_INFORMATION));

    //envPath is format "name=path", this is wrong
    //and we dont need it for createprocess
    gamePath = GetDirectory(gameExe);
    if(gamePath)
    {
        SetCurrentDirectoryA(gamePath);
    }

    #ifndef CNCNET_CLIENT
        config_init();

        #ifdef USE_UPNP
        upnp_run(config_port);
        #endif
    #endif

    int returnCode = CNCNET_ERROR_INVALID_PROCESS;

    snprintf(gameParamsFull, MAX_PATH, "%s %s", GetFile(gameExe), gameParams);

    // create process
    if(CreateProcessA(gameExe, (LPSTR)gameParamsFull, 0, 0, FALSE, DEBUG_PROCESS, 0, 0, &sInfo, &pInfo))
    {
        BOOL close = FALSE;
        
        while(!close)
        {
            #ifdef CNCNET_CLIENT__ // not active yet
            std::vector<Packet> data = LauncherHooks::outgoingPacketInstance()->getArrivedData();

            int i = 0;
            for(std::vector<Packet>::iterator iter = data.begin(); iter != data.end(); iter++)
            {
                Packet packet  = (*iter);
                // Send packets
            }
            #endif

            WaitForDebugEvent(&ev, INFINITE);

            DWORD dwContinue = DBG_CONTINUE;

            switch(ev.dwDebugEventCode)
            {
                case CREATE_PROCESS_DEBUG_EVENT:
                    processes++;

                    if(!hProcess)
                    {
                        game = games;
                        do
                        {
                            hk = game[0];

                            if(!hk->addr)
                            {
                                game++;
                                continue;
                            }

                            match = TRUE;
                            do
                            {
                                ReadProcessMemory(ev.u.CreateProcessInfo.hProcess, (char *)hk->addr, &byte, 1, NULL); //&dwWritten);

                                // look for CALL op
                                if(byte != 0xE8)
                                {
                                    match = FALSE;
                                }
                                hk++;
                            } while(hk->addr);

                            if(match)
                            {
                                dwProcessId = ev.dwProcessId;
                                hProcess = ev.u.CreateProcessInfo.hProcess;
                                hThread = ev.u.CreateProcessInfo.hThread;

                                hk = game[0];
                                do
                                {
                                    WriteProcessMemory(hProcess, (char *)hk->addr, &bp, sizeof(bp), &dwWritten);
                                    #if _DEBUG
                                    hk->cnt = 0;
                                    #endif
                                    hk++;
                                } while(hk->addr);

                                #if _DEBUG
                                fprintf(stderr, "launcher: Injected %d hooks\n", (hk - game[0]) );
                                #endif

                                // it is important to init wsock for us too, not only the debuggee
                                net_init();

                                SetProcessAffinityMask(hProcess, 0x01);
                                break;
                            }

                            game++;
                        } while(game[0]->addr);

                        #if _DEBUG
                        if(!match)
                        {
                            fprintf(stderr, "launcher: Warning: Unknown process %d!\n", processes);
                        }
                        #endif
                    }
                    #if _DEBUG
                    else
                    {
                        fprintf(stderr, "launcher: Notice: Ignoring new process %d, game already hooked\n", processes);
                    }
                    #endif

                    break;
                case EXIT_PROCESS_DEBUG_EVENT:

                    processes--;

                    if(ev.dwProcessId == dwProcessId && match)
                    {
                        
                        net_free();
                        #if _DEBUG
                        fprintf(stderr, "launcher: Used hooks:\n");
                        hk = game[0];
                        do
                        {
                            if(hk->cnt > 0)
                            {
                                fprintf(stderr, "launcher:  0x%08X: %d\n", (unsigned int)hk->addr, hk->cnt);
                            }
                            hk++;
                        } while(hk->addr);

                        fprintf(stderr, "launcher: Unused hooks:\n");
                        hk = game[0];
                        do
                        {
                            if(hk->cnt == 0)
                            {
                                fprintf(stderr, "launcher:  0x%08X\n", (unsigned int)hk->addr);
                            }
                            hk++;
                        } while(hk->addr);
                        #endif

                        returnCode = CNCNET_ERROR_NO_ERROR;
                        close = TRUE;
                    }

                    if(processes == 0 && !match)
                    {
                        close = TRUE;
                        returnCode = CNCNET_ERROR_INVALID_PROCESS;
                    }
                    #if _DEBUG
                    fprintf(stderr, "launcher: process terminated, %d processes left\n", processes);
                    #endif

                case OUTPUT_DEBUG_STRING_EVENT:
                    #if _DEBUG
                    ptr = (char *)malloc(ev.u.DebugString.nDebugStringLength);
                    ReadProcessMemory(hProcess, ev.u.DebugString.lpDebugStringData, ptr, ev.u.DebugString.nDebugStringLength, &dwWritten);
                    //fprintf(stderr, "launcher: game debug string: %s\n", ptr);
                    free(ptr);
                    #endif
                    break;

                case EXCEPTION_DEBUG_EVENT: 
                    if(ev.u.Exception.ExceptionRecord.ExceptionCode == EXCEPTION_ACCESS_VIOLATION)
                    {
                        dwContinue = DBG_EXCEPTION_NOT_HANDLED;
                    }
                    else if(ev.u.Exception.ExceptionRecord.ExceptionCode == EXCEPTION_BREAKPOINT)
                    {
                        memset(&ctx, 0, sizeof(CONTEXT));
                        ctx.ContextFlags = CONTEXT_CONTROL|CONTEXT_INTEGER;

                        GetThreadContext(hThread, &ctx);

                        hk = game[0];
                        do
                        {
                            if(hk->addr == (unsigned int)ev.u.Exception.ExceptionRecord.ExceptionAddress)
                            {
                                ((void(*)(CONTEXT *))hk->func)(&ctx);
                                #if _DEBUG
                                hk->cnt++;
                                #endif
                                break;
                            }
                            hk++;
                        } while(hk->addr);

                        SetThreadContext(hThread, &ctx);
                    }
                    else
                    {
                        dwContinue = DBG_EXCEPTION_NOT_HANDLED;
                    }
            } 

           ContinueDebugEvent(ev.dwProcessId, ev.dwThreadId, dwContinue);
        }
    }
    else
    {
        #if _DEBUG
        printf("CreateProcess failed with %d\n", (int)GetLastError());
        #endif
        return CNCNET_ERROR_FAILED_TO_LAUNCH;
    }

    Sleep(500);
    CloseHandle(pInfo.hProcess);
    CloseHandle(pInfo.hThread);

    return returnCode;
}

// hooks use these
uint32_t pop(CONTEXT *ctx)
{
    DWORD ret;
    DWORD tmp;

    ReadProcessMemory(hProcess, (void *)ctx->Esp, &ret, 4, &tmp);
    ctx->Esp += 4;

    return ret;
}

void push(CONTEXT *ctx, uint32_t val)
{
    DWORD tmp;
    WriteProcessMemory(hProcess, (void *)ctx->Esp, &val, 4, &tmp);
    ctx->Esp -= 4;
}

uint32_t memread(void *addr, void *buf, int len)
{
    DWORD tmp = 0;
    ReadProcessMemory(hProcess, addr, buf, len, &tmp);
    return tmp;
}

uint32_t memwrite(void *addr, void *buf, int len)
{
    DWORD tmp = 0;
    WriteProcessMemory(hProcess, addr, buf, len, &tmp);
    return tmp;
}
